import sys

import numpy as np
from collections import defaultdict

import torch
from torch.utils.data.distributed import DistributedSampler
import math
from torch.utils.data.sampler import Sampler


class NormalEpisodeSampler(Sampler):
    def __init__(self, cfg, *args, **kwargs):

        super().__init__(*args, **kwargs)

        self.cfg = cfg
        self.n_way = self.cfg.N_WAY
        self.n_shot = self.cfg.N_SHOT
        self.n_query = self.cfg.N_QUERY
        self.episode_size = self.cfg.EPISODE_SIZE
        self.labels = list(kwargs['data_source'].labels.numpy())

        # {0:[path1, path2, ...]}
        self.label2image = self._get_label2images()

        self.indices = self.update_indices()
        print('the episodes has been updated over !')

    def __len__(self):
        return self.episode_size

    def _get_label2images(self):
        # # image_indexs = np.arange(len(self.dataset.images))
        label2images = defaultdict(list)
        for index, label in enumerate(self.labels):
            label2images[label].append(index)
        return label2images

    def __iter__(self):

        return iter(self.indices)

    # 每经过一个epoch, 都应该手动调用该参数, 即sampler.indices = sampler.update_episode()
    def update_indices(self):
        indices = [self._sample_episode() for _ in range(self.episode_size)]
        return indices

    def _sample_episode(self):
        labels = list(set(self.labels))
        selected = np.random.choice(labels, size=self.n_way, replace=False)
        # the mixed set of support set and query set
        # support_query_set is a list contains n_way*(n_shot+n_query) images
        # [s1, s1, s1, s1, s1, s2, s2, ... , q1, q1, q1, q1, q1, q2, q2, ...]
        support_set = []
        query_set = []
        for i, c in enumerate(selected):
            # 因为path是str, 所以np会把第二个index元素也变成str
            candidates = self.label2image[c]
            indexs = np.random.choice(candidates, size=self.n_shot + self.n_query, replace=False)
            support_set.extend(indexs[:self.n_shot])
            query_set.extend(indexs[self.n_shot:])
            # labels: [0,1,...,n_way-1]
            # the first element of labels stands for the label of support1 and query1
        indices = support_set + query_set
        # indices = [int(x[1]) for x in support_query_set]

        return indices


class EpisodeSampler(DistributedSampler):
    def __init__(self, cfg, *args, **kwargs):

        super().__init__(*args, **kwargs)

        self.cfg = cfg
        self.n_way = self.cfg.N_WAY
        self.n_shot = self.cfg.N_SHOT
        self.n_query = self.cfg.N_QUERY
        self.episode_size = self.cfg.EPISODE_SIZE
        self.labels = list(self.dataset.labels.numpy())

        # {0:[path1, path2, ...]}
        self.label2image = self._get_label2images()

        # self.indices = torch.load('episode_indices.pth')
        # print('the episodes has been loaded over !')

        self.indices = self.update_indices()
        print('the episodes has been updated over !')

        # 如果设置drop_last = True
        self.num_samples = math.ceil(self.episode_size / self.num_replicas)

    def _get_label2images(self):
        # # image_indexs = np.arange(len(self.dataset.images))
        label2images = defaultdict(list)
        for index, label in enumerate(self.labels):
            label2images[label].append(index)
        return label2images

    def __iter__(self):
        indices = self.indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples, print(len(indices), self.num_samples)

        return iter(indices)

    # 每经过一个epoch, 都应该手动调用该参数, 即sampler.indices = sampler.update_episode()
    def update_indices(self):
        indices = [self._sample_episode() for _ in range(self.episode_size)]
        return indices

    def _sample_episode(self):
        labels = list(set(self.labels))
        selected = np.random.choice(labels, size=self.n_way, replace=False)
        # the mixed set of support set and query set
        # support_query_set is a list contains n_way*(n_shot+n_query) images
        # [s1, s1, s1, s1, s1, s2, s2, ... , q1, q1, q1, q1, q1, q2, q2, ...]
        support_set = []
        query_set = []
        for i, c in enumerate(selected):
            # 因为path是str, 所以np会把第二个index元素也变成str
            candidates = self.label2image[c]
            indexs = np.random.choice(candidates, size=self.n_shot + self.n_query, replace=False)
            support_set.extend(indexs[:self.n_shot])
            query_set.extend(indexs[self.n_shot:])
            # labels: [0,1,...,n_way-1]
            # the first element of labels stands for the label of support1 and query1
        indices = support_set + query_set
        # indices = [int(x[1]) for x in support_query_set]

        return indices
